# -*- coding: utf-8 -*-
# @Time    : 2020/4/26 上午9:53
# @Author  : tsl
# @Email   : mymailwith163@163.com
# @File    : learnQT.py
# @Software: PyCharm

from PyQt5 import QtCore, QtGui, QtWidgets

##########################################
import sys, os
import datetime,time
import cv2

from  tools import toolsClass

CurFolder = os.getcwd()
DefaultImFolder = CurFolder
NowTime = datetime.datetime.now()
Month = str(NowTime.month).zfill(2)
Day = str(NowTime.day).zfill(2)
Hour = str(NowTime.hour).zfill(2)
Minute = str(NowTime.minute).zfill(2)

class Ui_MainWindow(object):
    def setupUi(self, MainWindow):

        self.versionGet()
        self.version = self.versionGet()["version"]
        self.name = self.versionGet()["name"]
        MainWindow.setWindowTitle(self.name+"V"+self.version)

        self.totalTime = 0
        self.number = 0
        self.imgList = []
        self.sigleTime = []

        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(1100, 1300)

        font = QtGui.QFont()
        font.setPointSize(20)
        font.setBold(False)
        font.setWeight(75)
        MainWindow.setFont(font)

        self.window = QtWidgets.QWidget(MainWindow)
        self.window.setObjectName("window")

        self.horizontalLayout = QtWidgets.QHBoxLayout(self.window)
        self.horizontalLayout.setObjectName("horizontalLayout")
        self.horizontalLayout_2 = QtWidgets.QHBoxLayout()
        self.horizontalLayout_2.setObjectName("horizontalLayout_2")

        self.verticalLayout = QtWidgets.QVBoxLayout(self.window)
        self.verticalLayout.setObjectName("verticalLayout")

        self.OpenDirBnt = QtWidgets.QPushButton(self.window)
        self.OpenDirBnt.setObjectName("OpenDirBnt")
        self.OpenDirBnt.setFont(font)
        self.OpenDirBnt.setIconSize(QtCore.QSize(20,20))
        self.OpenDirBnt.setStyleSheet("background-color: rgb(255, 255, 255);")
        self.OpenDirBnt.setText("待检测目录")
        self.verticalLayout.addWidget(self.OpenDirBnt)

        self.PreImBnt = QtWidgets.QPushButton(self.window)
        self.PreImBnt.setObjectName("PreImBnt")
        self.PreImBnt.setFont(font)
        self.PreImBnt.setIconSize(QtCore.QSize(20,20))
        self.PreImBnt.setStyleSheet("background-color: rgb(255, 255, 255);")
        self.PreImBnt.setText("前一张图")
        self.verticalLayout.addWidget(self.PreImBnt)

        self.NextImBnt = QtWidgets.QPushButton(self.window)
        self.NextImBnt.setObjectName("PreImBnt")
        self.NextImBnt.setFont(font)
        self.NextImBnt.setIconSize(QtCore.QSize(20,20))
        self.NextImBnt.setStyleSheet("background-color: rgb(255, 255, 255);")
        self.NextImBnt.setText("后一张图")
        self.verticalLayout.addWidget(self.NextImBnt)

        self.UpdateBnt = QtWidgets.QPushButton(self.window)
        self.UpdateBnt.setObjectName("UpdateBnt")
        self.UpdateBnt.setFont(font)
        self.UpdateBnt.setIconSize(QtCore.QSize(20,20))
        self.UpdateBnt.setStyleSheet("background-color: rgb(255, 0, 0);")
        self.UpdateBnt.setText("升级")
        self.verticalLayout.addWidget(self.UpdateBnt)

        self.VersionBnt = QtWidgets.QPushButton(self.window)
        self.VersionBnt.setObjectName("VersionBnt")
        self.VersionBnt.setFont(font)
        self.VersionBnt.setIconSize(QtCore.QSize(20,20))
        self.VersionBnt.setStyleSheet("background-color: rgb(255, 0, 0);")
        self.VersionBnt.setText(self.version)
        self.verticalLayout.addWidget(self.VersionBnt)


        self.horizontalLayout.addLayout(self.verticalLayout)

        self.ImShowLabel = QtWidgets.QLabel(self.window)
        self.horizontalLayout.addWidget(self.ImShowLabel)
        self.horizontalLayout.setStretch(0, 1)
        self.horizontalLayout.setStretch(1, 8)

        self.ImShowLabel.setStyleSheet("background-color: rgb(0, 0, 0);")
        self.ImShowLabel.setText("当前图像")
        # self.ImShowLabel.setFrameRect(QtCore.QRect(0,0,500,500))

        self.ImShowLabel.setObjectName("ImShowLabel")

        self.horizontalLayout_2.addLayout(self.horizontalLayout)
        MainWindow.setCentralWidget(self.window)

        self.haveUpdate = False

        # self.menubar = QtWidgets.QMenuBar(MainWindow)
        # self.menubar.setGeometry(QtCore.QRect(0, 0, 1193, 26))
        # self.menubar.setObjectName("menubar")
        # MainWindow.setMenuBar(self.menubar)
        # self.statusbar = QtWidgets.QStatusBar(MainWindow)
        # self.statusbar.setObjectName("statusbar")
        # MainWindow.setStatusBar(self.statusbar)

        ################定义相关变量并初始化################
        self.ImFolder = ''  # 图片文件夹路径
        self.ImNameSet = []  # 图片集合
        self.CurImId = 0  # 当前显示图在集合中的编号
        self.MainWindow = MainWindow
        ################button按钮点击事件回调函数################
        self.OpenDirBnt.clicked.connect(self.OpenDirBntClicked)
        self.NextImBnt.clicked.connect(self.NextImBntClicked)
        self.PreImBnt.clicked.connect(self.PreImBntClicked)
        # self.UpdateBnt.clicked.connect(self.UpdateBntClicked)

    def toUpdateVersion(self):
        newVersion = 3.1
        version_name = self.versionGet()
        with open("version.txt","w") as f:
            f.writelines("version:"+str(newVersion)+"\n")
            f.writelines("name:"+version_name["name"])
        self.haveUpdate = True
        self.version = newVersion

    def versionGet(self):
        with open("version.txt","r")as f:
            lines = f.readlines()
            for line in lines:
                if(line.startswith("version")):
                    version = line.split(":")[-1].rstrip()
                if(line.startswith("name")):
                    name = line.split(":")[-1].rstrip()
        version = {"version":"版本V"+version,"name":name}
        return version

    def getTimeMsgs(self,ImFolder):
        totalTime = 0
        msg = {}
        for file in os.listdir(ImFolder):
            if(not self.haveUpdate):
                needTime = toolsClass.getTime()
            else:
                needTime = toolsClass.getTime() - toolsClass.lossTime()
            totalTime+=needTime
            # file_Time = {file:needTime}
            msg[os.path.join(ImFolder,file)]=needTime;
        msgs = {"img_and_time":msg,
                "total_time":totalTime}
        return msgs

    def cvtLoc(self,rx, ry, rw, rh, iw, ih):
        x2 = (rx + rw / 2.0) * iw
        x1 = (rx - rw / 2.0) * iw
        y2 = (ry + rh / 2.0) * ih
        y1 = (ry - rh / 2.0) * ih

        x, y, w, h = x1, y1, x2 - x1, y2 - y1
        return map(int, [x, y, w, h])

    def cvt(self,ImPath,msgs):

        totalTime = msgs["total_time"]
        img_and_time = msgs["img_and_time"]
        useTime = 0

        img = cv2.imread(ImPath)
        locPath = os.path.join("../labels", ImPath.split("/")[-2]) + "/" + ImPath.split("/")[-1].replace(".jpg", ".txt")
        if(os.path.exists(locPath)):
            name = os.path.splitext(locPath)[0]
            useTime = img_and_time[os.path.abspath(locPath.replace(".txt",".jpg")).replace("labels","dataset")]
            with open(locPath, "r") as f:
                _, iw, ih = img.shape[::-1]
                lines = f.readlines()
                for line in lines:
                    line = line.rstrip()
                    try:
                        label, x, y, w, h = map(float, line.split(" "))
                    except:
                        x, y, w, h = map(float, line.split(" ")[1::])
                        label = line.split(" ")[0]

                    rx, ry, rw, rh = self.cvtLoc(x, y, w, h, iw, ih)
                    try:
                        cv2.putText(img, str(int(float(label)))+",w="+str(rw)+",h="+str(rh)+",time="+str(useTime), (rx, ry), 1, 1, (0, 0, 255), 1)
                    except:
                        cv2.putText(img, str(label)+",w="+str(rw)+",h="+str(rh)+",time="+str(useTime), (rx, ry), 1, 1, (0, 0, 255), 1)

                    cv2.rectangle(img, (rx, ry), (rx + rw, ry + rh), (255, 0, 0), 2)

        img = cv2.resize(img, (1670, 1000))
        cv2.putText(img, "Total Time: " + str(totalTime) + "ms", (10, 60), 2, 2, (255, 255, 255), 2)
        cv2.putText(img, "This Image: " + str(useTime) + "ms", (10, 130), 2, 2, (255, 255, 255),2)
        if(os.path.exists(ImPath)):
            cv2.imwrite("tmp.jpg", img)
            ImPath = "tmp.jpg"

        return ImPath

    def OpenDirBntClicked(self):
        ImFolder = QtWidgets.QFileDialog.getExistingDirectory(None,"select folder", "../dataset")
        msgs = self.getTimeMsgs(ImFolder)

        if ImFolder != '':
            ImNameSet = []
            totalNum = len(os.listdir(ImFolder))
            self.imgList = os.listdir(ImFolder)

            self.totalTime = msgs["total_time"]
            time.sleep(self.totalTime/1000)


            for i in os.listdir(ImFolder):
                if(i.endswith("g")):
                    ImNameSet.append(i)

            ImNameSet.sort()
            ImPath = os.path.join(ImFolder, ImNameSet[0])

            ImPath = self.cvt(ImPath,msgs)

            pix = QtGui.QPixmap(ImPath)
            self.ImShowLabel.setPixmap(pix)

            self.ImFolder = ImFolder
            self.ImNameSet = ImNameSet
            self.CurImId = 0
            _, SelectFolderName = os.path.split(ImFolder)


            # _translate = QtCore.QCoreApplication.translate
            # CurWinTitle = "AI目标识别边缘计算盒-目标检测算法V"+str(self.version)+"                                                " + \
            #               "                                                             " + \
            #               SelectFolderName + '/' + ImNameSet[0]
            # self.MainWindow.setWindowTitle(_translate("MainWindow", CurWinTitle))
        else:
            print('请重新选择文件夹')

    def NextImBntClicked(self):
        try:
            print(self.number,self.imgList[self.number])
        except:
            print(self.number,self.imgList[self.number-1])
        if(self.number<len(self.imgList)):
            self.number+=1
            print(self.number,"-----------------------------")
        else:
            return

        ImFolder = self.ImFolder
        msgs = self.getTimeMsgs(ImFolder)
        ImNameSet = self.ImNameSet
        CurImId = self.CurImId
        ImNum = len(ImNameSet)
        if CurImId < ImNum - 1:  # 不可循环看图
            ImPath = os.path.join(ImFolder, ImNameSet[CurImId + 1])
            ImPath = self.cvt(ImPath,msgs)

            pix = QtGui.QPixmap(ImPath)

            self.ImShowLabel.setPixmap(pix)
            self.CurImId = CurImId + 1

    def PreImBntClicked(self):
        self.number -= 1
        if(self.number<0):
            self.number = 0

        print(self.number,self.imgList[self.number])
        ImFolder = self.ImFolder

        msgs = self.getTimeMsgs(ImFolder)

        ImNameSet = self.ImNameSet
        CurImId = self.CurImId
        ImNum = len(ImNameSet)
        if CurImId > 0:  # 第一张图片没有前一张
            ImPath = os.path.join(ImFolder, ImNameSet[CurImId - 1])
            # totalNum = len(os.listdir(ImFolder))
            ImPath = self.cvt(ImPath,msgs)

            pix = QtGui.QPixmap(ImPath)
            self.ImShowLabel.setPixmap(pix)
            self.CurImId = CurImId - 1

        if self.CurImId < 0:
            self.CurImId = 0




#########主函数入口 #########
if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    mainWindow = QtWidgets.QMainWindow()
    ui = Ui_MainWindow()
    ui.setupUi(mainWindow)
    mainWindow.show()
    sys.exit(app.exec_())

